from ray.rllib.models.torch.torch_modelv2 import TorchModelV2

from rl_evolve.models.dqn_model import DQNModel
from rl_evolve.models.ppo_model import PPOModel
from rl_evolve.models.sac_model import SACModel

MODEL_CLS_MAPPING = {
    "DQN": DQNModel,
    "PPO": PPOModel,
    "SAC": SACModel,
}


def get_model_class(algo_name: str) -> TorchModelV2:
    """Return the model class for the given algorithm name.

    None is returned if the algorithm name is not found in the mapping.
    """

    return MODEL_CLS_MAPPING.get(algo_name, None)
